import numpy as np
import tritonclient.http as httpclient
from transformers import BertTokenizer


class HttpClient:
    def __init__(self, modelname='prompt_bert', url='10.39.230.41:80', model_path='/data0/jianyu10/PTM/huggingface_model_cache/chinese-roberta-wwm-ext', verbose=False):
        self.model_name = modelname
        self.triton_client = httpclient.InferenceServerClient(url=url, verbose=verbose)
        self.tokenizer = BertTokenizer.from_pretrained(model_path)

    def gettmp(self, title):
        sentence = f'{title}，它的意思是[MASK]。'
        sen_tmp = f'{title}，这句话的意思是[MASK]。'
        return sentence, sen_tmp

    def getids(self, titles):
        batch_size = len(titles)
        inputs = []
        ori_inputs = []
        tmp_inputs = []
        inputs.append(httpclient.InferInput('input_ids', [batch_size, 123], "INT32"))
        inputs.append(httpclient.InferInput('input_tem', [batch_size, 123], "INT32"))
        for title in titles:
            sen, sen_tmp = self.gettmp(title)
            oriids = self.tokenizer.encode_plus(sen, max_length=123, padding='max_length')['input_ids']
            temids = self.tokenizer.encode_plus(sen_tmp, max_length=123, padding='max_length')['input_ids']
            ori_inputs.append(oriids)
            tmp_inputs.append(temids)

        input1 = np.array(ori_inputs, dtype=np.int32)
        input2 = np.array(tmp_inputs, dtype=np.int32)

        inputs[0].set_data_from_numpy(input1)
        inputs[1].set_data_from_numpy(input2)
        return inputs

    def inference(self, titles):
        """
        同步调用
        """
        inputs = self.getids(titles)
        outputs = [
            httpclient.InferRequestedOutput('output'),
        ]
        result = self.triton_client.infer(
            model_name=self.model_name,
            inputs=inputs,
            request_id=str('1'),
            outputs=outputs)
        output_array = result.as_numpy('output')
        return output_array

def test(pid):
    titles = ['英国首相府违规聚会调查报告认为政府“领导不力”---调查报告认为政府“领导不力”',
              '在欢声笑语中展现新时代新征程上精气神---2022年春节联欢晚会——在欢声笑语中展现新时代新征程上精气神',
              '因发生斗殴事件致2名犯人死亡 美国联邦监狱进入封锁状态---美国休斯敦发生枪击事件 致1名警员死亡',
              '英媒：泽连斯基下令三年扩军十万 敦促议员不要散布恐慌---泽连斯基签署法令：未来三年内扩军十万人',
              '冬季风暴渐平息 美部分地区降雪超50厘米---冬季风暴持续影响美国大部分地区',
              '英媒：特朗普大举筹款瞄准2024总统大选---特朗普又夸下海口：若2024再当选总统，将赦免国会大厦骚乱者',
              '多国政要和国际组织官员贺新春 祝新年如虎添翼---视频｜多位国际组织负责人及国家政要送上新春祝福',
              '“美国‘超额死亡’人数已近百万”---华尔街日报：疫情下，美国“超额死亡”人数已近百万',
              '日本外相：驻日美军入境新冠检测所用方法未被日本认可有效---外媒：驻日美军人员被曝离美前未进行新冠检测',
              '除夕，布林肯又发新闻公报拜年：愿虎年给所有人带来机遇、成功和健康---美国务卿布林肯拜年：愿虎年给所有人带来机遇、成功和健康']
    client = HttpClient()
    while 1:
        for i in titles:
            ori, sim = i.split('---')
            inputs = [ori for _ in range(5)]
            st = time.time()
            res = client.inference(inputs)
            ed = time.time()
            time.sleep(0.01)
            print(pid, len(inputs), (ed - st) * 1000)
if __name__ == '__main__':
    import os
    import time

    test(os.getpid())
